import os
os.environ['CUDA_VISIBLE_DEVICES'] = '5'

import gradio as gr
import torch
from transformers import AutoConfig, AutoModelForCausalLM
from janus.models import MultiModalityCausalLM, VLChatProcessor
from janus.utils.io import load_pil_images
from PIL import Image

import numpy as np
import os
import time
# import spaces  # Import spaces for ZeroGPU compatibility


# Load model and processor
# model_path = "deepseek-ai/Janus-Pro-7B"
model_path = "/mnt/nas/zhangshu/Janus-Pro-7B"
config = AutoConfig.from_pretrained(model_path)
language_config = config.language_config
language_config._attn_implementation = 'eager'
vl_gpt = AutoModelForCausalLM.from_pretrained(model_path,
                                             language_config=language_config,
                                             trust_remote_code=True)
if torch.cuda.is_available():
    vl_gpt = vl_gpt.to(torch.bfloat16).cuda()
else:
    vl_gpt = vl_gpt.to(torch.float16)

vl_chat_processor = VLChatProcessor.from_pretrained(model_path)
tokenizer = vl_chat_processor.tokenizer
cuda_device = 'cuda' if torch.cuda.is_available() else 'cpu'


@torch.inference_mode()
# @spaces.GPU(duration=120) 
# Multimodal Understanding function
def multimodal_understanding(image, question, seed=42, top_p=0.95, temperature=0.1):
    # Clear CUDA cache before generating
    torch.cuda.empty_cache()
    
    # set seed
    torch.manual_seed(seed)
    np.random.seed(seed)
    torch.cuda.manual_seed(seed)
    
    conversation = [
        {
            "role": "<|User|>",
            "content": f"<image_placeholder>\n{question}",
            "images": [image],
        },
        {"role": "<|Assistant|>", "content": ""},
    ]
    
    pil_images = [Image.fromarray(image)]
    prepare_inputs = vl_chat_processor(
        conversations=conversation, images=pil_images, force_batchify=True
    ).to(cuda_device, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float16)
    
    
    inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
    
    outputs = vl_gpt.language_model.generate(
        inputs_embeds=inputs_embeds,
        attention_mask=prepare_inputs.attention_mask,
        pad_token_id=tokenizer.eos_token_id,
        bos_token_id=tokenizer.bos_token_id,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=512,
        do_sample=False if temperature == 0 else True,
        use_cache=True,
        temperature=temperature,
        top_p=top_p,
    )
    
    answer = tokenizer.decode(outputs[0].cpu().tolist(), skip_special_tokens=True)
    return answer

from os.path import join as osj
dir = r'/mnt/nas/datasets/diction/ZipArchive0729'
save_clothing_dir = r'/mnt/nas/datasets/diction/ZipArchive0729_clo'
save_human_dir = r'/mnt/nas/datasets/diction/ZipArchive0729_human'
import os,shutil,cv2
filenames = os.listdir(dir)
if os.path.exists(save_clothing_dir):shutil.rmtree(save_clothing_dir)
os.makedirs(save_clothing_dir)
if os.path.exists(save_human_dir):shutil.rmtree(save_human_dir)
os.makedirs(save_human_dir)

Q1 = 'Is there anyone in the picture? Just answer yes or no.'  # no
Q2 = 'Is this the complete garment? Just answer yes or no.'   # yes
check_no = lambda ans : 'no' in ans.lower()
check_yes = lambda ans : 'yes' in ans.lower()

count = 0
for entry in os.scandir( dir ):
    if entry.is_file() and entry.name.endswith('.jpg'):
        filename = entry.name
    else:continue

    count += 1
    print('\rprocess num : ',count,end='',flush=True)    
    
    filepath = os.path.join( dir , filename)


    # results = model.predict(filepath)
    A1 = multimodal_understanding( cv2.imread( filepath ) , Q1 )
    A2 = multimodal_understanding( cv2.imread( filepath ) , Q2 )

    if check_no( A1 ) and check_yes( A2 ):
        shutil.copy2( osj(dir , filename) , osj(save_clothing_dir,filename) )
    else:
        shutil.copy2( osj(dir , filename) , osj(save_human_dir,filename) )
        